/*
 * To change this template, choose Tools | Templates
 * and open the template in the editor.
 */

package mygame.procedural.ik;

import com.jme3.animation.AnimControl;
import com.jme3.animation.Bone;
import com.jme3.animation.Skeleton;
import com.jme3.export.InputCapsule;
import com.jme3.export.JmeExporter;
import com.jme3.export.JmeImporter;
import com.jme3.export.OutputCapsule;
import com.jme3.math.FastMath;
import com.jme3.math.Quaternion;
import com.jme3.math.Transform;
import com.jme3.math.Vector3f;
import com.jme3.renderer.RenderManager;
import com.jme3.renderer.ViewPort;
import com.jme3.scene.Node;
import com.jme3.scene.Spatial;
import com.jme3.scene.control.AbstractControl;
import com.jme3.scene.control.Control;
import java.io.IOException;

/**
 *
 * @author Paulo
 */
public class Ik2JointControl extends AbstractControl {
    //Any local variables should be encapsulated by getters/setters so they
    //appear in the SDK properties window and can be edited.
    //Right-click a local variable to encapsulate it with getters and setters.

    private boolean refresh;
    private Skeleton skeleton;
    private Bone base;
    private float baseLength;
    private Bone tip;
    private float tipLength;
    private Bone ikTarget;
    private Node node;
    
    private Vector3f upDirection;
    private String ikTargetName;
    private String ikNodeName;
    private float angle;
    
    public Ik2JointControl() {
        refresh = false;
        upDirection = Vector3f.UNIT_Y;
        ikTargetName = null;
        ikNodeName = null;
        angle = 0.0f;
    }
    
    @Override
    protected void controlUpdate(float tpf) {
        refresh();
        
        if (spatial != null && node != null && base != null) {
            setUserControl(base, true);

            solveIk();

            setUserControl(base, false);
        }
    }

    @Override
    protected void controlRender(RenderManager rm, ViewPort vp) {
        //Only needed for rendering-related operations,
        //not called when spatial is culled.
    }

    @Override
    public Control cloneForSpatial(Spatial spatial) {
        Ik2JointControl control = new Ik2JointControl();
        control.setAngle(getAngle());
        control.setIkNodeName(getIkNodeName());
        control.setIkTargetName(getIkTargetName());
        control.setUpDirection(getUpDirection());
        
        return control;
    }

    @Override
    public void setSpatial(Spatial spatial) {
        super.setSpatial(spatial);
        
        if (spatial != null) {
            skeleton = spatial.getControl(AnimControl.class).getSkeleton();
            
            if (ikTargetName != null) {
                ikTarget = skeleton.getBone(getIkTargetName());
            
                tip = ikTarget.getParent();
                tipLength = ikTarget.getModelSpacePosition().distance(tip.getModelSpacePosition());

                base = tip.getParent();
                baseLength = base.getModelSpacePosition().distance(tip.getModelSpacePosition());
            }
            
            if (ikNodeName != null) {
                node = (Node) ((Node) spatial).getChild(getIkNodeName());
            }
            
        } else {
            skeleton = null;
            node = null;
            ikTarget = null;
            tip = null;
            tipLength = 0;
            base = null;
            baseLength = 0;
        }
    }
    
    @Override
    public void read(JmeImporter im) throws IOException {
        super.read(im);
        InputCapsule in = im.getCapsule(this);
        angle = in.readFloat("angle", 0.0f);
        ikNodeName = in.readString("ikNodeName", new String());
        ikTargetName = in.readString("ikTargetName", new String());
        upDirection = (Vector3f) in.readSavable("upDirection", Vector3f.UNIT_Y);
        refresh = true;
    }

    @Override
    public void write(JmeExporter ex) throws IOException {
        super.write(ex);
        OutputCapsule out = ex.getCapsule(this);
        out.write(angle, "angle", 0.0f);
        out.write(ikNodeName, "ikNodeName", new String());
        out.write(ikTargetName, "ikTargetName", new String());
        out.write(upDirection, "upDirection", Vector3f.UNIT_Y);
    }
    
    private void refresh() {
        if (refresh) {
            if (spatial != null) {
                skeleton = spatial.getControl(AnimControl.class).getSkeleton();
                
                if (skeleton != null && getIkTargetName() != null) {
                    ikTarget = skeleton.getBone(getIkTargetName());

                    tip = ikTarget.getParent();
                    tipLength = ikTarget.getModelSpacePosition().distance(tip.getModelSpacePosition());

                    base = tip.getParent();
                    baseLength = base.getModelSpacePosition().distance(tip.getModelSpacePosition());
                }
                
                if (getIkNodeName() != null) {
                    node = (Node) ((Node) spatial).getChild(getIkNodeName());
                }
            }
        }
        refresh = false;
    }
    
    private void solveIk() {
        Vector3f targetV = spatial.worldToLocal(node.getWorldTranslation(), null).subtract(base.getModelSpacePosition());
        
        // final local angles calculation
        float dist = targetV.length();
        float tipAngle = FastMath.acos((baseLength * baseLength + tipLength * tipLength - dist * dist) / (2 * baseLength * tipLength));
        float baseAngle = FastMath.acos((baseLength * baseLength + dist * dist - tipLength * tipLength) / (2 * baseLength * dist));
        
        Quaternion poleRotation = new Quaternion();
        poleRotation.fromAngleAxis(getAngle(), getUpDirection());
        Vector3f poleAxis = poleRotation.mult(Vector3f.UNIT_X);
        
        Quaternion lookAt = new Quaternion();
        lookAt.fromAngleAxis(getUpDirection().angleBetween(targetV.normalize()),
                getUpDirection().cross(targetV));
        
        Quaternion ikBase = new Quaternion();
        ikBase.fromAngleAxis(-baseAngle, poleAxis);

        ikBase.multLocal(lookAt);
        
        base.setUserTransformsWorld(base.getModelSpacePosition(), ikBase);
        
        Quaternion ikTip = new Quaternion();
        ikTip.fromAngleAxis(FastMath.PI - tipAngle - baseAngle, poleAxis);
        
        ikTip.multLocal(lookAt);
        
        tip.setUserTransformsWorld(base.getModelSpacePosition().add(tip.getModelSpacePosition()), ikTip);
        
        ikTarget.setUserTransformsWorld(ikTarget.getLocalPosition(), node.getLocalRotation());
        updateBoneTransformConservative(ikTarget);
        
        updateBoneTransform(base);
        
    }
    
    private void setUserControl(Bone bone, boolean value) {
        bone.setUserControl(value);
        for (Bone b : bone.getChildren()) {
            setUserControl(b, value);
        }
    }
    
    private void updateBoneTransform(Bone bone) {
        Transform t;
        for (Bone b : bone.getChildren()) {
            t = b.getCombinedTransform(bone.getModelSpacePosition(), bone.getModelSpaceRotation());
            b.setUserTransformsWorld(t.getTranslation(), b.getModelSpaceRotation());
            updateBoneTransform(b);
        }
    }
    private void updateBoneTransformConservative(Bone bone) {
        Transform t;
        for (Bone b : bone.getChildren()) {
            t = b.getCombinedTransform(bone.getModelSpacePosition(), bone.getModelSpaceRotation());
            b.setUserTransformsWorld(t.getTranslation(), t.getRotation());
            updateBoneTransformConservative(b);
        }
    }

    public Vector3f getUpDirection() {
        return upDirection;
    }

    public void setUpDirection(Vector3f upDirection) {
        this.upDirection = upDirection;
    }

    public String getIkTargetName() {
        return ikTargetName;
    }

    public void setIkTargetName(String ikTargetName) {
        this.ikTargetName = ikTargetName;
        refresh = true;
    }

    public String getIkNodeName() {
        return ikNodeName;
    }

    public void setIkNodeName(String ikNodeName) {
        this.ikNodeName = ikNodeName;
        refresh = true;
    }

    public float getAngle() {
        return angle;
    }

    public void setAngle(float angle) {
        this.angle = angle;
    }
}
